Repository navigation
[ROCm][DSv4] Fuse the shared expert into AITER MegaMoE - #61004
shen-shanshan merged 2 commits into
Conversation
With VLLM_ROCM_AITER_MEGA_MOE_FUSE_SHARED_EXPERTS=1, the aiter_mega_moe backend runs DeepSeek-V4's shared expert as one extra local expert per EP rank inside MegaMoEV2 instead of a separate FP8 MLP after it. Routed ids are remapped to the (routed_per_rank + 1) layout and every token gets a weight-1.0 route to its own rank's shared slot. The block-FP8 shared expert is requantized to MXFP4 into that slot when MegaMoE weights are finalized, and the MoRI heap default grows to 16G. Co-authored-by: Claude Signed-off-by: Yanfei Wang <yanfei.wang@amd.com>
mergify
Bot
added
deepseek
|
👋 Hi! Thank you for contributing to the vLLM project. 💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment Once the PR is approved or has the If you have any questions, please reach out to us on Slack at https://slack.vllm.ai. Agent GuidelinesIMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban. 🚀 |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
✅ @AMD-yanfeiwang, CI is now available for this PR.
|
|
/ci run |
|
/ci run |
|
✅ Triggered Buildkite CI #94035 for commit |
|
✅ CI is already running for this commit: https://buildkite.com/vllm/ci/builds/94035 |
|
@AMD-yanfeiwang Thanks for the work. |
@AMD-yanfeiwang please disclose NIAH accuracy and more comparisons. Requantizing shared experts just to fuse the shared experts has been known to cause large accuracy degradation, that is why AITER later developed a heterogeous fused shared expert MoE kernel instead of requantizing the shared experts. |
@benenzhu Thanks for the review. I added GPQA-Diamond and MMLU-Pro in thinking mode, with the official sampling settings, plus NIAH.
Full setup and per-eval details are in the PR description's Accuracy section. |
@tjtanaa Thanks. NIAH: I ran lm-eval's RULER needle tasks ( In thinking mode, GPQA-Diamond (0.894 → 0.899) and MMLU-Pro (0.852 → 0.854) are also unchanged; see the reply above. I see no accuracy cost from requantizing the shared expert on these evals.
In exchange, decode ITL drops 4–9% and throughput rises 3–8%. Full setup and per-eval details are in the PR description's Accuracy section. |
shen-shanshan
left a comment
shen-shanshan
left a comment
There was a problem hiding this comment.
The accuracy results look good to me.
CI selector (shadow): 97 test steps (162 jobs) instead of 74 (96 jobs)Shadow mode: this changes nothing about what CI runs. It shows what the evidence-based selector would pick for this PR, next to today's rules. How it works. Feedback welcome: reply here if it would skip a step this change needs, or runs something unrelated.
Selector would run (97)
Would skip (today's rules run them) (26)
Would add (today's rules do not run them) (49)
AMD mirrors: would skip (0)none AMD mirrors: would add (83)
5 changed files · base |
|
I don't think having Why do we have Users should expect the precision a model ships to be the one effectively used, and shoud use opt-ins if they want to override. AITER alreads hacks this with https://github.com/ROCm/aiter/blob/9fd025efd11c8d3edf660475d5b680fbb11cdc35/aiter/fused_moe.py#L777-L810 silently. If you want to use a different lower precision, please use Useful as well: #55686 |
7 participants
Overview
Adds an opt-in env var,
VLLM_ROCM_AITER_MEGA_MOE_FUSE_SHARED_EXPERTS=1. With it set,--moe-backend aiter_mega_moe(#59685) runs DeepSeek-V4's shared expert inside AITER MegaMoEV2, as one extra local expert per EP rank, instead of as a separate FP8 MLP after the MoE. This is the ROCm counterpart of the DeepGEMM MegaMoE shared-expert fusion (#53040), using the same extra-slot scheme as sglang on ROCm (sgl-project/sglang#35619).Claims
Performance
Setup
deepseek-ai/DeepSeek-V4-Pro.vllm/vllm-openai-rocm:nightly-7d0b4e57…, the merge base of this PR, with the PR files mounted.--data-parallel-size 8 --enable-expert-parallel --all2all-backend mori_high_throughput --moe-backend aiter_mega_moe--max-num-batched-tokens 8192 --kv-cache-dtype fp8 --async-scheduling --gpu-memory-utilization 0.86 --no-enable-prefix-cachingvllm bench serve --dataset-name random --ignore-eos --num-prompts max(16, 4×conc).ISL 1024 / OSL 1024
ISL 8192 / OSL 1024
Each point is a single run. In an earlier same-node session, two baseline runs differed by ≤0.6% in median ITL; with default flags that session also showed −4% to −7% ITL and +4% to +8% throughput at concurrency 1–128. TTFT is prefill/queue dominated and noisy at these prompt counts; the 1k/1k conc-128 TTFT increase did not reproduce at 8k/1k.
Accuracy
All runs use the AgentX DEP8 server flags from the Performance section; baseline and fused differ only in the env var.
Thinking mode, official sampling settings (vLLM's default
reasoning_effort=high,temperature=1.0, top_p=1.0,--reasoning-parser deepseek_v4, chat completions)Idavidrein/gpqa, 198 q, 0-shot, 128k output cap)Prompt: the simple-evals format, ending in
Answer: $LETTER; the score is theAnswer:line of the final content. Choices are shuffled with a fixed seed, so both arms see identical prompts.Arms: baseline and fused are within 1σ on both evals. GPQA-Diamond matches the official Think High score.
Truncation: responses still thinking at the output cap are scored wrong:
Excluding them, MMLU-Pro is 0.886 (baseline) vs 0.885 (fused), so the ~1.7-point gap to the official number is the output cap, not the fusion. The official Think Max setting recommends a ≥384K context.
Non-thinking / short-budget A/B (same node and image for both arms)
fingertap/GPQA-Diamondmirror of the same 198 questions.niah_single_{1,2,3},niah_multikey_{1,2,3},niah_multiquery,niah_multivalue), one run per context length.lm_eval --model local-completions --tasks gsm8k --num_fewshot 5. The difference is about 1.2σ; [ROCm][DSv4] AITER MegaMoEV2 Integration For DeepSeek V4 #59685 reported 0.9484 for the unfused MegaMoE path.DeepseekV4ForCausalLM, no vision encoder).Details
experts=392, topk=7. A small Triton kernel remaps routed ide → e + e // 48and appendsep_rank * 49 + 48with weight 1.0. The routed weights already includerouted_scaling_factor. The shared route always stays on the token's own rank, so it adds no all-to-all traffic.scaled_dequantize+mxfp4_quantize, i.e. aiterdynamic_mxfp4_quant) into the shared slot, and the MLP is dropped. MTP and DSpark layers go through the same finalize path.topk=7needs a larger MoRI heap, so its default rises from 8G to 16G when the flag is on.max_num_batched_tokens=16384, startup can hang in the profile run on MI355X, with or without this flag, on the unmodified image too.Synchronizerthat every stream shares. DSv4's CSA layers run the two compressortorch.mmcalls (N=2048 and N=512, fp32 out) concurrently on two aux streams with the same torch handle. At ≥12288 tokens per rank, the N=512 GEMM picks a Stream-K remainder kernel (…SK3_SKFTR0_SKXCCM8). One kernel can then clear a flag the other is waiting on, and that workgroup spins inlabel_SK_Fixupforever (seen with rocgdb). The other ranks then wait in the MegaMoE collective. It is not a MegaMoE bug.--max-num-batched-tokens 8192(the AgentX recipe). At that size the N=512 GEMM never takes the remainder path, and every run started reliably.Tests:
pytest tests/models/test_deepseek_v4_aiter_mega_moe.pypasses on MI355X (6 passed). The newtest_shared_expert_route_appendedcovers the id remap,-1padding routes, and the appended shared route. The end-to-end EP8 path needs 8× gfx950 and the DSv4 checkpoint, so it is validated outside CI (results above).Not a duplicate: the existing MegaMoE shared-fusion PRs are NVIDIA DeepGEMM (#53040, #56568, #59927) or Kimi K3 (#53556). #53847 targets the non-mega AITER fused-MoE path.
Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.
AI assistance (Claude) was used for the code, tests and this description. I reviewed every changed line and ran the validation above.